import torch

from torch import nn
from torch.utils.tensorboard import SummaryWriter

import dataset
from net import CIFARNet

class Trainer:
    def __init__(self, epoch_times: int, use_gpu: bool, learn_rate=0.01):
        self.loss = None
        self.epoch_no = 0
        self.train_times = 0
        self.epoch_times = epoch_times

        self.net = CIFARNet()
        self.loss_fn = nn.CrossEntropyLoss()  # 损失函数
        self.test_loader = dataset.get_loader(False)
        self.train_loader = dataset.get_loader(True)
        # https://ptorch.com/docs/1/optim
        # 保存当前的参数状态并且基于计算梯度更新参数
        self.optimizer = torch.optim.SGD(self.net.parameters(), lr=learn_rate)
        self.writer = SummaryWriter(log_dir='./logs')  # 使用tensorboard可视化数据

        if not use_gpu:
            self.device = None
        else:
            self.device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
            self.net = self.net.to(self.device)

    def train_record(self):
        if self.train_times % 100 == 0:
            print(f'训练次数：{self.train_times}，Loss:{self.loss.item()}')
            self.writer.add_scalar("train_loss", self.loss.item(), self.train_times)

    def train(self):
        self.train_times = 0
        for epoch_idx in range(self.epoch_times):
            self.epoch_no = epoch_idx
            self.net.train()
            print(f'-----第{self.epoch_no}轮训练开始-----')
            for data in self.train_loader:
                inputs, labels = data
                if self.device is not None:
                    inputs, labels = inputs.to(self.device), labels.to(self.device)
                self.optimizer.zero_grad()  # 优化器优化模型
                outputs = self.net(inputs)
                self.loss = self.loss_fn(outputs, labels)
                self.loss.backward()
                self.optimizer.step()
                self.train_times += 1
                self.train_record()
        self.writer.close()

    def test(self, epoch_idx: int):
        # net.eval()
        total_size = 0
        total_loss = 0
        total_accuracy = 0
        with torch.no_grad():
            for data in self.test_loader:
                images, labels = data
                if self.device is not None:
                    images, labels = images.to(self.device), labels.to(self.device)
                outputs = self.net(images)
                loss = self.loss_fn(outputs, labels)
                total_loss += loss.item()
                # argmax 找出每一张预测概率最大的下标
                # accuracy = (outputs.argmax(1) == labels).sum()
                _, predicted = torch.max(outputs.data, 1)
                total_size += labels.size(0)
                total_accuracy += (predicted == labels).sum().item()
        print(f'测试集的Loss:{total_loss}')
        print(f'测试集的正确率：{(total_accuracy / total_size)}')

    def save(self, path: str):
        torch.save(self.net, "cifdarnet.pth")

    def save_dict(self, path: str):
        torch.save({'epoch': self.epoch_no,
                    'model_state_dict': self.net.state_dict(),
                    'optimizer_state_dict': self.optimizer.state_dict(),
                    'loss': self.loss
                    }, path)
